# -*- coding: utf-8 -*-
# @author : wanglei
# @date : 2020/4/29 11:11 PM
# @description : 朴素贝叶斯，看病案例讲解
from manimlib.imports import *
from from_victor.common_scenes.custom_scenes import *
from from_victor.common_scenes.pluggable_scenes import *
from from_victor.common_widgets.custom_widgets import *

# 字幕
captions = [
    "看病时，医生会通过你的症状，根据经验给出初步诊断",
    "那症状和经验在数学里怎么表示呢？",
    "我们来看概率论中的贝叶斯公式，别见公式就怕，很简单",
    "上面就是贝叶斯公式，咱们用得新冠肺炎和具有发热症状来理解下",
    "这里将logo与np.array(LEFT*3)左对齐、右对齐",
    "还可以传入coor_mask参数(三维ndarray)，来决定在哪个维度上进行移动",
    "比如，move_to(UP*2)，但是将y轴的移动屏蔽掉，就是这样的效果"
]

"""
标题
"""
class Title(TitleScene):
    CONFIG = {
        'title': '数学如何看病?'
    }
    def setup(self):
        logo_gp = Logo_gp()
        logo_gp.scale(0.7)
        logo_gp.move_to(3 * LEFT + 6 * UP)
        self.play(FadeInFromLarge(logo_gp), run_time=2)
        super().setup()
        self.wait()

"""
引出问题，医生看病的过程
"""
class Issue(Scene):
    CONFIG={

    }

    def construct(self):
        logo_gp = Logo_gp()
        logo_gp.scale(0.7)
        logo_gp.move_to(3 * LEFT + 6 * UP)
        self.add(logo_gp)
        man = SVGMobject('man.svg').set_stroke(width=0).set_color('#6545A4')
        man.move_to(2 * LEFT)
        man_hot = man.copy()
        man_hot.set_color('#ff2d51') # 红色
        self.play(GrowFromPoint(man, 3 * LEFT))
        symptom = TextMobject('症状\\\发\\\热\\\.\\\. ', color=BLACK).scale(0.7)
        symptom.next_to(man,RIGHT,aligned_edge=UP)
        self.play(Write(symptom))
        self.play(Transform(man, man_hot))
        self.play(WiggleOutThenIn(man_hot))
        arrow = Arrow(start=ORIGIN, end= 3 * RIGHT, color='#6545A4', stroke_width=3)
        arrow.next_to(symptom,RIGHT)
        self.play(GrowFromCenter(arrow))
        text1 = TextMobject('医生', color = BLACK).scale(0.5)
        text2 = TextMobject('初步诊断', color = BLACK).scale(0.5)
        text1.next_to(arrow, UP)
        text2.next_to(arrow, DOWN)
        self.play(FadeInFromPoint(text1,text1.get_center() + UP), FadeInFromDown(text2))
        text3 = TextMobject('感\\\冒', color = BLACK).scale(0.7)
        text3.next_to(arrow, RIGHT)
        self.play(GrowFromPoint(text3, 4 * RIGHT))
        self.wait()


"""
讲解贝叶斯公式，和它的变形-朴素贝叶斯公式
"""
class NaiveBayes(Scene):
    def construct(self):
        logo_gp = Logo_gp()
        logo_gp.scale(0.7)
        logo_gp.move_to(3 * LEFT + 6 * UP)
        self.play(Write(logo_gp))

        start_len = len(self.get_mobjects())
        # 画概率图
        A = Circle(arc_center=1 * LEFT, radius=1.3, stroke_width=0, fill_color=RED, fill_opacity=0.8)
        B = Circle(arc_center=1 * RIGHT, radius=1.3, stroke_width=0, fill_color=BLUE, fill_opacity=0.8)
        self.play(ShowCreation(A), ShowCreation(B))

        self.play(WiggleOutThenIn(A))
        A_text = TextMobject('A', color = BLACK).scale(0.7)
        A_text.move_to(A.get_center())
        self.play(FadeIn(A_text))
        A_line = Line(A_text.get_top(), A_text.get_top() + 1.5 * UP, color = RED)
        A_line.next_to(A_text, UP)
        self.play(Write(A_line))
        A_text_info = TextMobject('A事件发生的概率', color=BLACK).scale(0.5)
        A_text_info.next_to(A_line, UP)
        self.play(FadeIn(A_text_info))

        self.play(WiggleOutThenIn(B))
        B_text = TextMobject('B', color=BLACK).scale(0.7)
        B_text.move_to(B.get_center())
        self.play(FadeIn(B_text))
        B_line = Line(B_text.get_top(), B_text.get_top() + 2 * UP, color=BLUE)
        B_line.next_to(B_text, UP)
        self.play(Write(B_line))
        B_text_info = TextMobject('B事件发生的概率', color=BLACK).scale(0.5)
        B_text_info.next_to(B_line, UP)
        self.play(FadeIn(B_text_info))

        # 弧线用来圈出指定区域
        p1, p2 = circle_intersection(A, B)
        arc1 = ArcBetweenPoints(start=p1, end=p2, angle=0.44 * PI, color=GREEN, stroke_width=3)
        arc2 = ArcBetweenPoints(start=p1, end=p2, angle=- 0.44 * PI, color=GREEN, stroke_width=3)
        gp1 = VGroup(arc1, arc2)
        self.play(WiggleOutThenIn(gp1))
        AB_text = TexMobject('{A}\\\\\cap\\\{B}', color=BLACK).scale(0.5)
        AB_text.move_to(gp1.get_center())
        self.play(FadeIn(AB_text))
        AB_line = Line(AB_text.get_top(), AB_text.get_top() + 0.5 * UP, color=GREEN)
        AB_line.next_to(AB_text, UP)
        self.play(Write(AB_line))
        AB_text_info = TextMobject('A和B同时发生的概率', color=BLACK).scale(0.5)
        AB_text_info.next_to(AB_line, UP)
        self.play(FadeIn(AB_text_info))

        line_path = Line(start=ORIGIN, end=3 * UP)
        graph_gp= VGroup(*(self.get_mobjects()[start_len:]))
        self.play(MoveAlongPath(graph_gp, line_path), rate_func=linear)


        # 条件概率公式
        pb1 = TexMobject('P(A|B) = \\frac{P({A}\\cap{B})}{P(B)}', color=BLACK)
        pb1[0][7:13].set_color('#6545A4')
        pb1.scale(0.7)
        self.play(Write(pb1))

        pb2 = TexMobject('P(B|A) = \\frac{P({A}\\cap{B})}{P(A)}', color=BLACK)
        pb2[0][7:13].set_color('#6545A4')
        pb2.scale(0.7)
        pb2.next_to(pb1, DOWN)
        self.play(Write(pb2))

        sub1 = pb1[0][7:13].copy()
        self.play(ApplyMethod(sub1.next_to, pb2, 2 * LEFT + DOWN))
        sub2 = pb1[0][6].copy()
        self.play(ApplyMethod(sub2.next_to, sub1, RIGHT))
        sub3 = pb1[0][:6].copy()
        self.play(ApplyMethod(sub3.next_to, sub2, RIGHT))
        sub4 = pb1[0][14:].copy()
        self.play(ApplyMethod(sub4.next_to, sub3, RIGHT))

        sub5 = pb2[0][6].copy()
        self.play(ApplyMethod(sub5.next_to, sub4, RIGHT))
        sub6 = pb2[0][:6].copy()
        self.play(ApplyMethod(sub6.next_to, sub5, RIGHT))
        sub7 = pb2[0][14:].copy()
        self.play(ApplyMethod(sub7.next_to, sub6, RIGHT))

        # 朴素贝叶斯公式
        nb = TexMobject('\\text{得出：}P(A|B) = \\frac{{P(B|A)}\\times{P(A)}}{P(B)}', color=BLACK)
        nb.next_to(sub3, DOWN)
        nb.scale(0.7)
        self.play(Write(nb))
        self.wait(12)

"""
求解感冒和新冠肺炎概率
"""
class Answer(Scene):
    def construct(self):
        logo_gp = Logo_gp()
        logo_gp.scale(0.7)
        logo_gp.move_to(3 * LEFT + 6 * UP)
        self.play(SpinInFromNothing(logo_gp))
        # 朴素贝叶斯公式
        nb = TexMobject('\\text{朴素贝叶斯公式：}P(A|B) = \\frac{{P(B|A)}\\times{P(A)}}{P(B)}', color=BLACK)
        nb.move_to(4.5 * UP)
        nb.scale(0.7)
        self.add(nb)

        # 计算感冒和新冠肺炎的概率
        start_len = len(self.get_mobjects())
        man = SVGMobject('man.svg').set_stroke(width=0).set_color('#6545A4')
        man.move_to(2.5 * LEFT + 3 * UP)
        self.play(GrowFromPoint(man, 3 * LEFT))
        symptom = TextMobject('症状：\\\发热\\\呼吸困难\\\不打喷嚏', color=BLACK).scale(0.7)
        symptom[0][:3].set_color('#6545A4')
        symptom.next_to(man, RIGHT, aligned_edge=UP)
        self.play(Write(symptom))
        arrow = Arrow(start=ORIGIN, end=3 * RIGHT, color='#6545A4', stroke_width=3)
        arrow.next_to(symptom, RIGHT)
        self.play(GrowFromCenter(arrow))
        text1 = TextMobject('朴素贝叶斯', color=BLACK).scale(0.5)
        text2 = TextMobject('初步诊断', color=BLACK).scale(0.5)
        text1.next_to(arrow, UP)
        text2.next_to(arrow, DOWN)
        self.play(FadeInFromPoint(text1, text1.get_center() + UP), FadeInFromDown(text2))
        text3 = Text('？', color=RED, font='smmyingbi').scale(1.5)
        text3.next_to(arrow, RIGHT)
        self.play(Write(text3))
        self.play(ShowPassingFlashAround(text3, run_time=2))
        pre_gp = VGroup(*(self.get_mobjects()[start_len:]))
        self.play(FadeOutAndShiftDown(pre_gp))

        p_new = TextMobject('P(新冠肺炎|发热,呼吸困难,不打喷嚏) = '
                            '\\\P(新冠肺炎|发热)'
                            '\\\P(新冠肺炎|呼吸困难)'
                            '\\\P(新冠肺炎|不打喷嚏)', color=BLACK).scale(0.7)
        p_new.next_to(nb, DOWN)
        p_cold = TextMobject('P(普通感冒|发热,呼吸困难,不打喷嚏) = '
                             '\\\P(普通感冒|发热)'
                             '\\\P(普通感冒|呼吸困难)'
                             '\\\P(普通感冒|不打喷嚏)', color=BLACK).scale(0.7)
        p_cold.next_to(p_new, DOWN)
        self.play(FadeInFromDown(p_new))
        self.play(FadeInFromDown(p_cold))
        self.wait(1)

        nb_actual = TexMobject('P(\\text{新冠肺炎}|\\text{发热}) = \\qquad\\qquad\\qquad \\\ '
                               '\\frac{{P(\\text{发热}|\\text{新冠肺炎})}\\times{P(\\text{新冠肺炎})}}{P(\\text{发热})}', color=BLACK)
        nb_actual.next_to(p_cold, DOWN)
        nb_actual.scale(0.7)
        self.play(Write(nb_actual))
        self.wait(6)


class Test(Scene):
    def construct(self):
        line1 = Line(3*LEFT, RIGHT, color=BLACK) #TODO：线的颜色会变粗，不知道咋回事；如果设置stroke_width为1，画完后颜色竟然还会变浅，也不知道咋回事
        self.play(Write(line1))
        t = TextMobject('你好中国', color=BLACK)
        t.next_to(line1, DOWN)
        self.play(Write(t))

class Test2(Scene):
    def construct(self):
        # if you want to color "{" or "}" braces character separately
        numeros4 = TexMobject( "\\mathbb{N}=\\{1,2,3,\\ldots\\}")
        # coloring first 2 characters
        numeros4[0][0:2].set_color(GREEN)
        # coloring 3rd character
        numeros4[0][2].set_color(RED)
        self.add(numeros4)
        self.wait()